package com.shujia.mllib

import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object Demo6TrainImage {
  def main(args: Array[String]): Unit = {

    /**
      * 训练模型
      *
      */

    val spark: SparkSession = SparkSession.builder()
      .appName("kmeans")
      .master("local[8]")
      .config("spark.sql.shuffle.partitions", "2")
      .getOrCreate()

    import spark.implicits._
    import org.apache.spark.sql.functions._


    //1、 读取数据
    val data: DataFrame = spark
      .read
      .format("libsvm")
      .load("spark/data/images")

    //将数据拆分成训练集和测试集
    val split: Array[Dataset[Row]] = data.randomSplit(Array(0.7, 0.3))
    val train: Dataset[Row] = split(0)
    val text: Dataset[Row] = split(1)


    //构建算法
    val regression = new LogisticRegression()
    regression.setFitIntercept(true)
    regression.setMaxIter(100)


    //将数据带入算法训练模型
    val model: LogisticRegressionModel = regression.fit(train)


    //将测试集带入模型判断模型准确率
    val dataFrame: DataFrame = model.transform(text)


    //计算准确率
    val p: DataFrame = dataFrame.select(sum(when($"label" === $"prediction", 1).otherwise(0)) / count($"label") as "p")

    p.show()


    //保存模型
    model
      .write
      .overwrite()
      .save("spark/data/imagemodel")


  }
}
